import streamlit as st
import streamlit.components.v1 as components
import os
import logging
from datetime import datetime
import base64
from pathlib import Path

# 导入自定义模块
try:
    from agent.langchain_agent import geographic_map_agent
    agent_available = True
except ImportError as e:
    st.error(f"导入智能代理失败: {e}")
    geographic_map_agent = None
    agent_available = False

try:
    from database.connection import db_manager
    db_available = True
except ImportError as e:
    st.warning(f"数据库模块不可用: {e}")
    db_manager = None
    db_available = False

try:
    from config import MAP_CONFIG
except ImportError as e:
    st.warning(f"配置文件加载失败: {e}")
    MAP_CONFIG = {
        'default_center': [39.9042, 116.4074],  # 北京
        'default_zoom': 10,
        'output_dir': 'outputs',
        'temp_dir': 'temp'
    }

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

def setup_page():
    """设置页面配置"""
    st.set_page_config(
        page_title="智能绘图助手",
        page_icon="🎨",
        layout="wide",
        initial_sidebar_state="expanded"
    )
    
    # 自定义CSS样式
    st.markdown("""
    <style>
    .main {
        padding-top: 2rem;
    }
    .stTextInput > div > div > input {
        background-color: #f0f2f6;
    }
    .chat-message {
        padding: 1rem;
        border-radius: 0.5rem;
        margin-bottom: 1rem;
        display: flex;
        flex-direction: column;
    }
    .user-message {
        background-color: #e3f2fd;
        border-left: 4px solid #2196f3;
    }
    .assistant-message {
        background-color: #f1f8e9;
        border-left: 4px solid #4caf50;
    }
    .image-container {
        border: 2px solid #ddd;
        border-radius: 10px;
        padding: 1rem;
        background-color: #fafafa;
        min-height: 600px;
        text-align: center;
    }
    .sidebar .sidebar-content {
        background-color: #f8f9fa;
    }
    </style>
    """, unsafe_allow_html=True)

def initialize_session_state():
    """初始化会话状态"""
    if "messages" not in st.session_state:
        st.session_state.messages = []
    if "current_image" not in st.session_state:
        st.session_state.current_image = None
        
        # 检查是否有现有的图片文件
        output_dir = Path(MAP_CONFIG['output_dir'])
        if output_dir.exists():
            # 查找图片文件 (PNG和GIF)
            image_files = list(output_dir.glob("*.png")) + list(output_dir.glob("*.gif"))
            if image_files:
                # 使用第一个找到的图片
                st.session_state.current_image = str(image_files[0])
                logger.info(f"找到图片文件: {image_files[0]}")
    
    if "processing" not in st.session_state:
        st.session_state.processing = False

def display_chat_message(message_type, content):
    """显示聊天消息"""
    css_class = "user-message" if message_type == "user" else "assistant-message"
    icon = "👤" if message_type == "user" else "🤖"
    
    st.markdown(f"""
    <div class="chat-message {css_class}">
        <div style="display: flex; align-items: center; margin-bottom: 0.5rem;">
            <span style="font-size: 1.2rem; margin-right: 0.5rem;">{icon}</span>
            <strong>{'用户' if message_type == 'user' else '助手'}</strong>
        </div>
        <div>{content}</div>
    </div>
    """, unsafe_allow_html=True)

def display_image(image_path):
    """显示图片结果"""
    if not image_path:
        st.info("暂无绘图结果")
        return
        
    try:
        # 确保路径是绝对路径
        if not os.path.isabs(image_path):
            image_path = os.path.abspath(image_path)
            
        if not os.path.exists(image_path):
            st.error(f"图片文件不存在: {image_path}")
            return
            
        if image_path.endswith(('.png', '.jpg', '.jpeg')):
            # 显示静态图片
            st.image(image_path, use_container_width=True)
            st.success(f"✅ 静态图片已加载: {os.path.basename(image_path)}")
            
        elif image_path.endswith('.gif'):
            # 显示GIF动画
            with open(image_path, 'rb') as f:
                gif_data = f.read()
            st.image(gif_data)
            st.success(f"✅ 动画图片已加载: {os.path.basename(image_path)}")
            
        else:
            st.warning(f"不支持的图片文件格式: {os.path.splitext(image_path)[1]}")
            
    except Exception as e:
        st.error(f"显示图片时出错: {str(e)}")
        logger.error(f"显示图片失败 - 路径: {image_path}, 错误: {e}")
        
        # 显示调试信息
        with st.expander("调试信息"):
            st.write(f"文件路径: {image_path}")
            st.write(f"文件存在: {os.path.exists(image_path) if image_path else False}")
            st.write(f"文件大小: {os.path.getsize(image_path) if image_path and os.path.exists(image_path) else 'N/A'}")
            st.write(f"错误详情: {str(e)}")

def sidebar_controls():
    """侧边栏控制面板"""
    with st.sidebar:
        st.header("🛠️ 控制面板")
        
        # 系统状态
        st.subheader("📊 系统状态")
        if agent_available:
            st.success("✅ 智能代理正常")
        else:
            st.error("❌ 智能代理不可用")
            
        if db_available:
            st.success("✅ 数据库模块正常")
            # 数据库连接测试
            if st.button("测试数据库连接"):
                with st.spinner("测试连接中..."):
                    try:
                        if db_manager and hasattr(db_manager, 'test_connection'):
                            is_connected = db_manager.test_connection()
                            if is_connected:
                                st.success("✅ 数据库连接正常")
                            else:
                                st.error("❌ 数据库连接失败")
                        else:
                            st.error("❌ 数据库管理器不可用")
                    except Exception as e:
                        st.error(f"❌ 连接测试出错: {str(e)}")
        else:
            st.error("❌ 数据库模块不可用")
        
        # 绘图配置
        st.subheader("🎨 绘图配置")
        default_center = st.text_input(
            "默认中心点 (纬度,经度)", 
            value=f"{MAP_CONFIG['default_center'][0]},{MAP_CONFIG['default_center'][1]}"
        )
        
        # 输出目录信息
        st.subheader("📁 输出信息")
        st.info(f"输出目录: {MAP_CONFIG['output_dir']}")
        st.info(f"临时目录: {MAP_CONFIG['temp_dir']}")
        
        # 图片文件管理
        st.subheader("🖼️ 图片文件")
        output_dir = Path(MAP_CONFIG['output_dir'])
        if output_dir.exists():
            # 获取所有支持的图片文件 (PNG和GIF)
            image_extensions = ['*.png', '*.gif']
            image_files = []
            for ext in image_extensions:
                image_files.extend(output_dir.glob(ext))
            
            if image_files:
                st.write(f"找到 {len(image_files)} 个图片文件:")
                
                # 创建文件选择器
                file_names = [f.name for f in image_files]
                selected_file = st.selectbox("选择要显示的图片:", file_names, 
                                           index=0 if st.session_state.current_image else 0)
                
                if st.button("🖼️ 加载选中图片"):
                    output_path = MAP_CONFIG.get('output_dir', 'outputs')
                    if output_path and selected_file:
                        selected_path = os.path.join(str(output_path), str(selected_file))
                        st.session_state.current_image = selected_path
                        st.success(f"已选择图片: {selected_file}")
                        st.rerun()
                    else:
                        st.error("输出目录或文件配置错误")
                
                # 显示当前图片信息
                if st.session_state.current_image:
                    current_file = Path(st.session_state.current_image).name
                    st.info(f"当前显示: {current_file}")
            else:
                st.warning("暂无图片文件")
                if st.button("🎯 生成演示图片"):
                    with st.spinner("生成演示图片中..."):
                        try:
                            import subprocess
                            result = subprocess.run(['python', 'demo_data.py'], 
                                                  capture_output=True, text=True)
                            if result.returncode == 0:
                                st.success("演示图片生成成功！")
                                st.rerun()
                            else:
                                st.error(f"生成失败: {result.stderr}")
                        except Exception as e:
                            st.error(f"生成演示图片失败: {e}")
        else:
            st.warning("输出目录不存在")
            if st.button("📁 创建输出目录"):
                output_dir.mkdir(parents=True, exist_ok=True)
                st.success("输出目录已创建")
                st.rerun()
        
        # 清空对话历史
        if st.button("🗑️ 清空对话历史"):
            st.session_state.messages = []
            st.session_state.current_image = None
            if geographic_map_agent and hasattr(geographic_map_agent, 'clear_context'):
                geographic_map_agent.clear_context()
            st.success("对话历史已清空")
            st.rerun()
        
        # 系统信息
        st.subheader("ℹ️ 系统信息")
        st.info("版本: 1.0.0")
        st.info("开发者: 智能绘图助手")
        st.info(f"当前时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")

def main_chat_interface():
    """主要聊天界面"""
    st.header("🎨 智能绘图助手")
    st.markdown("---")
    
    # 创建两列布局
    image_col, chat_col = st.columns([1.2, 1])
    
    with image_col:
        with st.container():
            # st.markdown('<div class="image-container">', unsafe_allow_html=True)
            
            if st.session_state.current_image:
                display_image(st.session_state.current_image)
            else:
                st.markdown("""
                <div style="text-align: center; padding: 100px 0; color: #666;">
                    <h3>🎨 绘图结果将在这里显示</h3>
                    <p>请在右侧聊天框中输入您的绘图需求</p>
                    <p>支持生成PNG静态图片和GIF动画</p>
                    <p>或点击快捷示例按钮快速开始</p>
                </div>
                """, unsafe_allow_html=True)
            
            st.markdown('</div>', unsafe_allow_html=True)
    
    with chat_col:
        st.subheader("💬 对话区域")
        
        # 显示欢迎信息（仅在没有消息时显示）
        if not st.session_state.messages:
            display_chat_message("assistant", """🎨 欢迎使用智能绘图助手！

我可以帮您：
• 📊 创建各种类型的数据可视化图表  
• 📈 制作统计图和分析图
• 🎨 生成美观的静态图片(PNG)
• 🎬 制作动态图片动画(GIF)
• 🗺️ 绘制地理数据相关的图表

请告诉我您想要绘制什么样的图，或者点击下方的快捷示例开始体验！""")
        
        # 显示聊天历史
        chat_container = st.container()
        with chat_container:
            for message in st.session_state.messages:
                display_chat_message(message["role"], message["content"])
        
        # 快捷示例按钮
        st.markdown("### 🚀 快捷示例")
        example_col1, example_col2 = st.columns(2)
        
        # 定义变量存储选中的示例
        selected_example = None
        
        with example_col1:
            if st.button("📊 统计图表", key="btn_chart"):
                selected_example = "创建一个包含随机数据的柱状图，保存为PNG格式"
            if st.button("📈 散点图", key="btn_scatter"):
                selected_example = "生成一个散点图，显示两个变量之间的关系，保存为PNG"
        
        with example_col2:
            if st.button("🎬 动画图表", key="btn_animation"):
                selected_example = "创建一个显示数据变化的动画图表，保存为GIF格式"
            if st.button("🌍 地理数据图", key="btn_geo"):
                selected_example = "绘制一个简单的地理数据图表，保存为PNG"
        
        # 用户输入
        st.markdown("---")
        user_input = st.text_input(
            "请输入您的绘图需求：",
            placeholder="例如：创建一个显示销售数据的柱状图",
            disabled=st.session_state.processing,
            key="user_input"
        )
        
        col1, col2, col3 = st.columns([1, 1, 2])
        with col1:
            send_button = st.button("发送", disabled=st.session_state.processing or not user_input, key="btn_send")
        with col2:
            if st.session_state.processing:
                st.spinner("处理中...")
        
        # 处理用户输入或快捷示例
        request_to_process = None
        
        if send_button and user_input:
            request_to_process = user_input
        elif selected_example:
            request_to_process = selected_example
        
        if request_to_process:
            # 添加用户消息
            st.session_state.messages.append({"role": "user", "content": request_to_process})
            st.session_state.processing = True
            
            # 处理请求
            with st.spinner("正在处理您的请求..."):
                try:
                    if agent_available and geographic_map_agent:
                        response = geographic_map_agent.process_request(request_to_process)
                        
                        # 检查是否生成了新图片
                        output_dir = Path(MAP_CONFIG['output_dir'])
                        if output_dir.exists():
                            # 获取所有支持的图片文件
                            image_extensions = ['*.png', '*.gif']
                            image_files = []
                            for ext in image_extensions:
                                image_files.extend(output_dir.glob(ext))
                            
                            if image_files:
                                # 获取最新的图片文件
                                try:
                                    latest_image = max(image_files, key=lambda f: f.stat().st_mtime)
                                    st.session_state.current_image = str(latest_image)
                                    logger.info(f"检测到新图片文件: {latest_image}")
                                except (OSError, ValueError) as e:
                                    logger.warning(f"无法确定最新图片文件: {e}")
                                    # 使用第一个找到的文件
                                    if image_files:
                                        st.session_state.current_image = str(image_files[0])
                    else:
                        response = """❌ 智能代理暂时不可用，请检查系统配置和依赖安装情况。

🔧 可能的解决方案：
1. 检查Python环境和依赖包安装
2. 确认配置文件中的API密钥设置
3. 验证数据库连接配置
4. 重启应用程序

如果问题持续存在，请联系系统管理员。"""
                    
                    # 添加助手回复
                    st.session_state.messages.append({"role": "assistant", "content": response})
                    
                except Exception as e:
                    error_msg = f"""❌ 处理请求时出现错误：

**错误详情：** {str(e)}

🔧 **建议操作：**
1. 检查网络连接状态
2. 确认输入格式正确
3. 重试您的请求
4. 如问题持续，请联系技术支持

💡 **提示：** 您可以尝试使用更简单的请求来测试系统功能。"""
                    st.session_state.messages.append({"role": "assistant", "content": error_msg})
                    logger.error(f"处理用户请求失败: {e}")
            
            st.session_state.processing = False
            # 清空输入框
            if 'user_input' in st.session_state:
                del st.session_state['user_input']
            st.rerun()

def main():
    """主函数"""
    try:
        # 设置页面
        setup_page()
        
        # 初始化会话状态
        initialize_session_state()
        
        # 显示侧边栏控制面板
        sidebar_controls()
        
        # 显示主界面
        main_chat_interface()
        
    except Exception as e:
        st.error(f"应用启动失败: {str(e)}")
        logger.error(f"应用启动失败: {e}")
        
        # 显示错误详情和建议
        with st.expander("查看错误详情"):
            st.code(str(e))
            
        st.markdown("""
        ### 🔧 解决建议：
        1. 检查所有依赖是否正确安装：`pip install -r requirements.txt`
        2. 确认配置文件存在且设置正确
        3. 验证输出目录权限
        4. 重启应用程序
        """)

if __name__ == "__main__":
    main() 